import threading
import time
from PyQt5.QtCore import pyqtSignal, QObject
from PyQt5.QtGui import QPixmap, QImage
import cv2
import torch
from ultralytics import YOLO  # 确保已 pip install ultralytics

class DetectionThread(threading.Thread, QObject):
    frame_signal = pyqtSignal(QPixmap, QPixmap, float)  # 原图, 检测图, FPS
    error_signal = pyqtSignal(str)

    def __init__(self, source, weights, conf=0.5, iou=0.5, device="cpu"):
        threading.Thread.__init__(self)
        QObject.__init__(self)
        self.source = source
        self.weights = weights
        self.conf = conf
        self.iou = iou
        self.device = device

        self._running = True
        self.paused = False

    def run(self):
        try:
            cap = cv2.VideoCapture(int(self.source) if self.source.isdigit() else self.source)
            model = YOLO(self.weights)
            model.fuse()  # optional 加速

            prev_time = time.time()
            while self._running and cap.isOpened():
                if self.paused:
                    time.sleep(0.1)
                    continue

                ret, frame = cap.read()
                if not ret:
                    break

                results = model.predict(frame, conf=self.conf, iou=self.iou, verbose=False)
                det_frame = results[0].plot()  # 返回带框图像

                # 转换为 QPixmap
                orig_pixmap = self.cv2pixmap(frame)
                det_pixmap = self.cv2pixmap(det_frame)

                fps = 1.0 / (time.time() - prev_time + 1e-6)
                prev_time = time.time()

                self.frame_signal.emit(orig_pixmap, det_pixmap, fps)
        except Exception as e:
            self.error_signal.emit(str(e))
        finally:
            cap.release()

    def stop(self):
        self._running = False

    def pause(self):
        self.paused = True

    def resume(self):
        self.paused = False

    def cv2pixmap(self, frame):
        """BGR OpenCV -> QPixmap"""
        rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        qimg = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        return QPixmap.fromImage(qimg)
